use super::parser::*;

pub fn ast_print(program: &Program) {
    let mut formatter = AstFormatter::new(program);
    formatter.format();
    print!("{}", formatter.buff);
}

struct AstFormatter<'a> {
    program: &'a Program,
    buff: String,
    tab_num: u32,
}

impl<'a> AstFormatter<'a> {
    fn new(program: &'a Program) -> Self {
        AstFormatter {
            program,
            buff: String::new(),
            tab_num: 0,
        }
    }

    fn format(&mut self) {
        self.visit_program(self.program);
    }

    fn tab_print(&mut self) {
        for _ in 0..self.tab_num {
            self.buff.push_str("    ");
        }
    }

    fn tab_entry(&mut self) {
        self.tab_num += 1;
    }

    fn tab_exit(&mut self) {
        self.tab_num -= 1;
    }

    fn visit_program(&mut self, program: &Program) {
        for comp_unit in &program.comp_unit_list {
            self.visit_comp_unit(comp_unit);
        }
    }

    fn visit_comp_unit(&mut self, comp_unit: &CompUnit) {
        match comp_unit {
            CompUnit::Decl(decl) => self.visit_decl(decl),
            CompUnit::FuncDef(func_def) => self.visit_func_def(func_def),
        }
    }

    fn visit_decl(&mut self, decl: &Decl) {
        match decl {
            Decl::ConstDecl(const_decl) => self.visit_const_decl(const_decl),
            Decl::VarDecl(var_decl) => self.visit_var_decl(var_decl),
        }
    }

    fn visit_const_decl(&mut self, const_decl: &ConstDecl) {
        self.buff.push_str("const");
        self.buff.push(' ');
        self.visit_b_type(&const_decl.b_type);
        self.buff.push(' ');
        for const_def in &const_decl.const_def_list {
            self.visit_const_def(const_def);
        }
        self.buff.push(';');
        self.buff.push('\n');
    }

    fn visit_b_type(&mut self, b_type: &BType) {
        match b_type {
            BType::Int => self.buff.push_str("int"),
        }
    }

    fn visit_const_def(&mut self, const_def: &ConstDef) {
        self.buff.push_str(&const_def.ident.clone());
        for const_exp in &const_def.const_exp_list {
            self.buff.push('[');
            self.visit_const_exp(const_exp);
            self.buff.push(']');
        }
        self.buff.push(' ');
        self.buff.push('=');
        self.buff.push(' ');
        self.visit_const_init_val(&const_def.const_init_val);
    }

    fn visit_const_init_val(&mut self, const_init_val: &ConstInitVal) {
        match &const_init_val {
            ConstInitVal::ConstExp(const_exp) => self.visit_const_exp(const_exp),
            ConstInitVal::ConstInitValList(const_init_val_list) => {
                self.buff.push('{');
                for (i, const_init_val) in const_init_val_list.iter().enumerate() {
                    self.visit_const_init_val(const_init_val);
                    if i != const_init_val_list.len() - 1 {
                        self.buff.push(',');
                        self.buff.push(' ');
                    }
                }
                self.buff.push('}');
            }
        }
    }

    fn visit_var_decl(&mut self, var_decl: &VarDecl) {
        self.visit_b_type(&var_decl.b_type);
        self.buff.push(' ');
        for (i, var_def) in var_decl.var_def_list.iter().enumerate() {
            self.visit_var_def(var_def);
            if i != var_decl.var_def_list.len() - 1 {
                self.buff.push(',');
                self.buff.push(' ');
            }
        }
        self.buff.push(';');
        self.buff.push('\n');
    }

    fn visit_var_def(&mut self, var_def: &VarDef) {
        self.buff.push_str(&var_def.ident.clone());
        for const_exp in &var_def.const_exp_list {
            self.buff.push('[');
            self.visit_const_exp(const_exp);
            self.buff.push(']');
        }
        if let Some(init_val) = &var_def.init_val {
            self.buff.push(' ');
            self.buff.push('=');
            self.buff.push(' ');
            self.visit_init_val(init_val);
        }
    }

    fn visit_init_val(&mut self, init_val: &InitVal) {
        match &init_val {
            InitVal::Exp(exp) => self.visit_exp(exp),
            InitVal::InitValList(init_val_list) => {
                self.buff.push('{');
                for (i, init_val) in init_val_list.iter().enumerate() {
                    self.visit_init_val(init_val);
                    if i != init_val_list.len() - 1 {
                        self.buff.push(',');
                        self.buff.push(' ');
                    }
                }
                self.buff.push('}');
            }
        }
    }

    fn visit_func_def(&mut self, func_def: &FuncDef) {
        if !self.buff.is_empty() {
            self.buff.push('\n');
        }
        self.visit_func_type(&func_def.func_type);
        self.buff.push(' ');
        self.buff.push_str(&func_def.ident.clone());
        self.buff.push('(');
        if let Some(func_f_params) = &func_def.func_f_params {
            self.visit_func_f_params(func_f_params);
        }
        self.buff.push(')');
        self.buff.push(' ');
        self.visit_block(&func_def.block);
    }

    fn visit_func_type(&mut self, func_type: &FuncType) {
        match func_type {
            FuncType::Int => self.buff.push_str("int"),
            FuncType::Void => self.buff.push_str("void"),
        }
    }

    fn visit_func_f_params(&mut self, func_f_params: &FuncFParams) {
        for (i, func_f_param) in func_f_params.func_f_param_list.iter().enumerate() {
            self.visit_func_f_param(func_f_param);
            if i != func_f_params.func_f_param_list.len() - 1 {
                self.buff.push(',');
                self.buff.push(' ');
            }
        }
    }

    fn visit_func_f_param(&mut self, func_f_param: &FuncFParam) {
        self.visit_b_type(&func_f_param.b_type);
        self.buff.push(' ');
        self.buff.push_str(&func_f_param.ident.clone());
        if let Some(exp_list) = &func_f_param.exp_list {
            self.buff.push('[');
            self.buff.push(']');
            for exp in exp_list {
                self.buff.push('[');
                self.visit_exp(exp);
                self.buff.push(']');
            }
        }
    }

    fn visit_block(&mut self, block: &Block) {
        self.buff.push('{');
        self.tab_entry();
        self.buff.push('\n');
        for block_item in &block.block_item_list {
            self.visit_block_item(block_item);
        }
        self.tab_exit();
        self.tab_print();
        self.buff.push('}');
        self.buff.push('\n');
    }

    fn visit_block_item(&mut self, block_item: &BlockItem) {
        self.tab_print();
        match block_item {
            BlockItem::Decl(decl) => self.visit_decl(decl),
            BlockItem::Stmt(stmt) => self.visit_stmt(stmt),
        }
    }

    fn visit_stmt(&mut self, stmt: &Stmt) {
        match stmt {
            Stmt::Assign(l_val, exp) => {
                self.visit_l_val(l_val);
                self.buff.push(' ');
                self.buff.push('=');
                self.buff.push(' ');
                self.visit_exp(exp);
                self.buff.push(';');
                self.buff.push('\n');
            }
            Stmt::Exp(exp) => {
                if let Some(exp) = exp {
                    self.visit_exp(exp);
                }
                self.buff.push(';');
                self.buff.push('\n');
            }
            Stmt::Block(block) => {
                self.visit_block(block);
            }
            Stmt::If(cond, stmt, else_stmt) => {
                self.buff.push_str("if");
                self.buff.push(' ');
                self.buff.push('(');
                self.visit_cond(cond);
                self.buff.push(')');
                if let Stmt::Block(_) = stmt.as_ref() {
                    self.buff.push(' ');
                    self.visit_stmt(stmt);
                } else {
                    self.tab_entry();
                    self.buff.push('\n');
                    self.tab_print();
                    self.visit_stmt(stmt);
                    self.tab_exit();
                }
                if let Some(else_stmt) = else_stmt {
                    self.tab_print();
                    self.buff.push_str("else");
                    if let Stmt::If(_, _, _) = else_stmt.as_ref() {
                        self.buff.push(' ');
                    } else if let Stmt::Block(_) = else_stmt.as_ref() {
                        self.buff.push(' ');
                    } else {
                        self.tab_entry();
                        self.buff.push('\n');
                        self.tab_print();
                        self.tab_exit();
                    }
                    self.visit_stmt(else_stmt);
                }
            }
            Stmt::While(cond, stmt) => {
                self.buff.push_str("while");
                self.buff.push(' ');
                self.buff.push('(');
                self.visit_cond(cond);
                self.buff.push(')');
                if let Stmt::Block(_) = stmt.as_ref() {
                    self.buff.push(' ');
                    self.visit_stmt(stmt);
                } else {
                    self.tab_entry();
                    self.buff.push('\n');
                    self.tab_print();
                    self.visit_stmt(stmt);
                    self.tab_exit();
                }
            }
            Stmt::Break => {
                self.buff.push_str("break;");
                self.buff.push('\n');
            }
            Stmt::Continue => {
                self.buff.push_str("continue;");
                self.buff.push('\n');
            }
            Stmt::Return(exp) => {
                self.buff.push_str("return");
                if let Some(exp) = exp {
                    self.buff.push(' ');
                    self.visit_exp(exp);
                }
                self.buff.push(';');
                self.buff.push('\n');
            }
        }
    }

    fn visit_exp_op(&mut self, exp_op: &ExpOp) {
        match &exp_op {
            ExpOp::Add => self.buff.push('+'),
            ExpOp::Sub => self.buff.push('-'),
            ExpOp::Mul => self.buff.push('*'),
            ExpOp::Div => self.buff.push('/'),
            ExpOp::Mod => self.buff.push('%'),
        }
    }

    fn visit_exp(&mut self, exp: &Exp) {
        match &exp {
            Exp::Primary(primary_exp) => {
                self.visit_primary_exp(primary_exp);
            }
            Exp::Inner(l_exp, op, r_exp) => {
                self.visit_exp(l_exp);
                self.buff.push(' ');
                self.visit_exp_op(op);
                self.buff.push(' ');
                self.visit_exp(r_exp);
            }
            Exp::Unary(op, exp) => {
                self.visit_unary_op(op);
                self.visit_exp(exp);
            }
        }
    }

    fn visit_cond_op(&mut self, cond_op: &CondOp) {
        match &cond_op {
            CondOp::Or => self.buff.push_str("||"),
            CondOp::And => self.buff.push_str("&&"),
            CondOp::Eq => self.buff.push_str("=="),
            CondOp::Neq => self.buff.push_str("!="),
            CondOp::Lt => self.buff.push('<'),
            CondOp::Gt => self.buff.push('>'),
            CondOp::Le => self.buff.push_str("<="),
            CondOp::Ge => self.buff.push_str(">="),
        }
    }

    fn visit_cond(&mut self, cond: &Cond) {
        match &cond {
            Cond::Exp(exp) => self.visit_exp(exp),
            Cond::CondExp(r_cond, op, l_cond) => {
                self.visit_cond(r_cond);
                self.buff.push(' ');
                self.visit_cond_op(op);
                self.buff.push(' ');
                self.visit_cond(l_cond);
            }
        }
    }

    fn visit_l_val(&mut self, l_val: &LVal) {
        self.buff.push_str(&l_val.ident.clone());
        for exp in &l_val.exp_list {
            self.buff.push('[');
            self.visit_exp(exp);
            self.buff.push(']');
        }
    }

    fn visit_unary_op(&mut self, unary_op: &UnaryOp) {
        match &unary_op {
            UnaryOp::Add => self.buff.push('+'),
            UnaryOp::Sub => self.buff.push('-'),
            UnaryOp::Not => self.buff.push('!'),
        }
    }

    fn visit_primary_exp(&mut self, primary_exp: &PrimaryExp) {
        match &primary_exp {
            PrimaryExp::Number(_num, str) => self.buff.push_str(str),
            PrimaryExp::FuncRParams(name, func_r_params) => {
                self.buff.push_str(name);
                self.buff.push('(');
                if let Some(func_r_params) = func_r_params {
                    self.visit_func_r_params(func_r_params);
                }
                self.buff.push(')');
            }
            PrimaryExp::LVal(l_val) => self.visit_l_val(l_val),
            PrimaryExp::Exp(exp) => {
                self.buff.push('(');
                self.visit_exp(exp);
                self.buff.push(')');
            }
        }
    }

    fn visit_func_r_params(&mut self, func_r_params: &FuncRParams) {
        for (i, func_f_param) in func_r_params.exp_list.iter().enumerate() {
            self.visit_exp(func_f_param);
            if i != func_r_params.exp_list.len() - 1 {
                self.buff.push(',');
                self.buff.push(' ');
            }
        }
    }

    fn visit_const_exp(&mut self, const_exp: &ConstExp) {
        self.visit_exp(&const_exp.exp);
    }
}
